#ifdef __APPLE__

#include "mach_hook.h"

#include <dlfcn.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>

#ifndef SEG_DATA_CONST
#define SEG_DATA_CONST  "__DATA_CONST"
#endif

struct rebindings_entry {
	struct rebinding *rebindings;
	size_t rebindings_nel;
	struct rebindings_entry *next;
};

static int prepend_rebindings(struct rebindings_entry **rebindings_head,
		struct rebinding rebindings[], size_t nel) {
	struct rebindings_entry *new_entry = (struct rebindings_entry *) malloc(
			sizeof(struct rebindings_entry));
	if (!new_entry) {
		return -1;
	}
	new_entry->rebindings = (struct rebinding *) malloc(
			sizeof(struct rebinding) * nel);
	if (!new_entry->rebindings) {
		free(new_entry);
		return -1;
	}
	memcpy(new_entry->rebindings, rebindings, sizeof(struct rebinding) * nel);
	new_entry->rebindings_nel = nel;
	new_entry->next = *rebindings_head;
	*rebindings_head = new_entry;
	return 0;
}

static void perform_rebinding_with_section(struct rebindings_entry *rebindings,
		section_t *section, intptr_t slide, nlist_t *symtab, char *strtab,
		uint32_t *indirect_symtab) {
	uint32_t *indirect_symbol_indices = indirect_symtab + section->reserved1;
	void **indirect_symbol_bindings = (void **) ((uintptr_t) slide
			+ section->addr);
	for (uint i = 0; i < section->size / sizeof(void *); i++) {
		uint32_t symtab_index = indirect_symbol_indices[i];
		if (symtab_index == INDIRECT_SYMBOL_ABS
				|| symtab_index == INDIRECT_SYMBOL_LOCAL
				|| symtab_index
						== (INDIRECT_SYMBOL_LOCAL | INDIRECT_SYMBOL_ABS)) {
			continue;
		}
		uint32_t strtab_offset = symtab[symtab_index].n_un.n_strx;
		char *symbol_name = strtab + strtab_offset;
		bool symbol_name_longer_than_1 = symbol_name[0] && symbol_name[1];
		struct rebindings_entry *cur = rebindings;
		while (cur) {
			for (uint j = 0; j < cur->rebindings_nel; j++) {
				if (symbol_name_longer_than_1
						&& strcmp(&symbol_name[1], cur->rebindings[j].name)
								== 0) {
					if (cur->rebindings[j].replaced != NULL
							&& indirect_symbol_bindings[i]
									!= cur->rebindings[j].replacement) {
						*(cur->rebindings[j].replaced) =
								indirect_symbol_bindings[i];
					}
					indirect_symbol_bindings[i] =
							cur->rebindings[j].replacement;
					goto symbol_loop;
				}
			}
			cur = cur->next;
		}
		symbol_loop: ;
	}
}

static void rebind_symbols_for_image(struct rebindings_entry *rebindings,
		const struct mach_header *header, intptr_t slide) {
	Dl_info info;
	if (dladdr(header, &info) == 0) {
		return;
	}

	segment_command_t *cur_seg_cmd;
	segment_command_t *linkedit_segment = NULL;
	struct symtab_command* symtab_cmd = NULL;
	struct dysymtab_command* dysymtab_cmd = NULL;

	uintptr_t cur = (uintptr_t) header + sizeof(mach_header_t);
	for (uint i = 0; i < header->ncmds; i++, cur += cur_seg_cmd->cmdsize) {
		cur_seg_cmd = (segment_command_t *) cur;
		if (cur_seg_cmd->cmd == LC_SEGMENT_ARCH_DEPENDENT) {
			if (strcmp(cur_seg_cmd->segname, SEG_LINKEDIT) == 0) {
				linkedit_segment = cur_seg_cmd;
			}
		} else if (cur_seg_cmd->cmd == LC_SYMTAB) {
			symtab_cmd = (struct symtab_command*) cur_seg_cmd;
		} else if (cur_seg_cmd->cmd == LC_DYSYMTAB) {
			dysymtab_cmd = (struct dysymtab_command*) cur_seg_cmd;
		}
	}

	if (!symtab_cmd || !dysymtab_cmd || !linkedit_segment
			|| !dysymtab_cmd->nindirectsyms) {
		return;
	}

	// Find base symbol/string table addresses
	uintptr_t linkedit_base = (uintptr_t) slide + linkedit_segment->vmaddr
			- linkedit_segment->fileoff;
	nlist_t *symtab = (nlist_t *) (linkedit_base + symtab_cmd->symoff);
	char *strtab = (char *) (linkedit_base + symtab_cmd->stroff);

	// Get indirect symbol table (array of uint32_t indices into symbol table)
	uint32_t *indirect_symtab = (uint32_t *) (linkedit_base
			+ dysymtab_cmd->indirectsymoff);

	cur = (uintptr_t) header + sizeof(mach_header_t);
	for (uint i = 0; i < header->ncmds; i++, cur += cur_seg_cmd->cmdsize) {
		cur_seg_cmd = (segment_command_t *) cur;
		if (cur_seg_cmd->cmd == LC_SEGMENT_ARCH_DEPENDENT) {
			if (strcmp(cur_seg_cmd->segname, SEG_DATA) != 0
					&& strcmp(cur_seg_cmd->segname, SEG_DATA_CONST) != 0) {
				continue;
			}
			for (uint j = 0; j < cur_seg_cmd->nsects; j++) {
				section_t *sect =
						(section_t *) (cur + sizeof(segment_command_t)) + j;
				if ((sect->flags & SECTION_TYPE) == S_LAZY_SYMBOL_POINTERS) {
					perform_rebinding_with_section(rebindings, sect, slide,
							symtab, strtab, indirect_symtab);
				}
				if ((sect->flags & SECTION_TYPE) == S_NON_LAZY_SYMBOL_POINTERS) {
					perform_rebinding_with_section(rebindings, sect, slide,
							symtab, strtab, indirect_symtab);
				}
			}
		}
	}
}

int rebind_symbols_image(void *header, intptr_t slide,
		struct rebinding rebindings[], size_t rebindings_nel) {
	struct rebindings_entry *rebindings_head = NULL;
	int retval = prepend_rebindings(&rebindings_head, rebindings,
			rebindings_nel);
	rebind_symbols_for_image(rebindings_head,
			(const struct mach_header *) header, slide);
	if (rebindings_head) {
		free(rebindings_head->rebindings);
	}
	free(rebindings_head);
	return retval;
}

#endif //!__APPLE__
